import json
import aiohttp


class Msg:
    role = "user"
    content = ""

    def __init__(self, role, content):
        self.role = role
        self.content = content

    def __str__(self):
        # 返回对象的可打印字符串表示
        return json.dumps({
            'role': self.role,
            'content': self.content
        })


async def chat(messages):
    url = 'https://cnb.cool/ai/chat/completions'
    headers = {
        'accept': 'text/event-stream',
        'accept-language': 'zh-CN,zh;q=0.9',
        'content-type': 'text/plain;charset=UTF-8',
        'origin': 'https://cnb.cool',
        'priority': 'u=1, i',
        'referer': 'https://cnb.cool/',
        'sec-ch-ua': '"Google Chrome";v="131", "Chromium";v="131", "Not_A Brand";v="24"',
        'sec-ch-ua-mobile': '?0',
        'sec-ch-ua-platform': 'Windows',
        'sec-fetch-dest': 'empty',
        'sec-fetch-mode': 'cors',
        'sec-fetch-site': 'same-origin',
        'user-agent': 'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/131.0.0.0 Safari/537.36'
    }

    messages2 = []
    data = {
        'model': 'hunyuan-chat',
        'stream': True,
        'messages': messages2
    }

    for x in messages:
        messages2.append(x.__dict__)
    # [{"role":"user","content":"\n1"},
    # {"role":"assistant","content":"您好，您似乎输入了一个数字“1”，需要我帮您做什么？请您提供更多的信息或者问一个具体的问题，我会尽力帮助您。"},
    # {"role":"user","content":"1"}]

    print(data)

    async for event_data in sse_client(url, headers, data):
        if event_data:
            if event_data == '[DONE]':
                return
            try:
                event_data = event_data.lstrip("data: ")
                # 尝试将事件数据解析为JSON对象
                event = json.loads(event_data)
                yield event["choices"][0]["delta"]["content"]
            except json.JSONDecodeError:
                # 如果解析失败，打印错误信息并继续处理下一个事件
                print(f"Failed to decode JSON: {event_data}")


async def sse_client(url, headers, json):
    async with aiohttp.ClientSession() as session:
        async with session.post(url, headers=headers, json=json) as response:
            # 确保响应类型是text/event-stream
            print(response.content_type)
            assert response.content_type == 'text/event-stream'

            async for line in response.content:
                decoded_line = line.decode('utf-8').strip()
                if decoded_line.startswith('data: '):
                    data = decoded_line[6:].strip()
                    # 在这里处理接收到的数据
                    yield data
